import numpy as np
import json
from typing import List, Dict
from collections import defaultdict

from matplotlib import pyplot as plt
import seaborn as sns


class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        return super().default(obj)

# 在TrafficGenerator中添加节点映射方法
class TrafficGenerator: # 流量生成器
    def __init__(self,
                 num_nodes: int = 4, # 定义节点数量
                 # min_value: float = 500,
                 # max_value: float = 800,
                 burst_config: Dict = None):
        self.num_nodes = num_nodes
        self.tor_nodes = [f"ToR{i+1}" for i in range(num_nodes)]  # 生成ToR名称列表（从1到4）
        self.time = 0.0
        # self.burst_intensity = 800

        # 突发流量配置
        self.burst_config = burst_config or {
            'hotspots': [0, 3],  # 突发热点节点
            'burst_duration': 0.5,  # 突发持续时间(s)
            'lambda_poisson': 0.3,  # 泊松分布参数
            # 'burst_intensity': 800.0  # 突发流量强度
        }

    def generate_all_to_all(self,
                            min_value: float = 500,
                            max_value: float = 800) -> np.ndarray:
        """
        生成可配置范围的All-to-All流量矩阵

        参数:
            min_value (float): 最小流量值，默认500GB
            max_value (float): 最大流量值，默认800GB

        返回:
            np.ndarray: N x N的流量矩阵，对角线置零
        """
        # 参数有效性校验
        if min_value < 0 or max_value < 0:
            raise ValueError("流量值不能为负数")
        if min_value >= max_value:
            raise ValueError("最大流量值必须大于最小流量值")

        # 生成均匀分布流量矩阵
        matrix = np.random.uniform(min_value, max_value,
                                   (self.num_nodes, self.num_nodes))
        np.fill_diagonal(matrix, 0)  # 对角线置零
        return matrix

    def generate_ring_allreduce(self,
                            min_value: float = 500,
                            max_value: float = 800) -> np.ndarray:
        """生成环形AllReduce流量矩阵"""
        matrix = np.zeros((self.num_nodes, self.num_nodes))
        chunk_size = np.random.uniform(min_value, max_value) # 定义每个节点的流量大小

        # 环形通信模式
        for i in range(self.num_nodes):
            # 顺时针方向通信
            src = i
            dst = (i + 1) % self.num_nodes
            matrix[src, dst] = chunk_size

            # 逆时针方向通信
            src = (i - 1) % self.num_nodes
            dst = i
            matrix[src, dst] = chunk_size

        return matrix

    def generate_burst_traffic(self, burst_intensity: float = 800) -> np.ndarray:
        """生成突发流量矩阵"""
        matrix = np.zeros((self.num_nodes, self.num_nodes))

        # 突发热点通信
        for src in self.burst_config['hotspots']:
            for dst in range(self.num_nodes):
                if src != dst:
                    matrix[src][dst] += np.random.exponential(
                        # self.burst_config['burst_intensity']
                        burst_intensity
                    )

        # 泊松过程控制突发次数
        if np.random.poisson(self.burst_config['lambda_poisson']) > 0:
            hotspot = np.random.choice(self.burst_config['hotspots'])
            for _ in range(int(np.random.exponential(5))):
                dst = np.random.choice(self.num_nodes)
                if hotspot != dst:
                    matrix[hotspot][dst] += np.random.exponential(
                        self.burst_config['burst_intensity'] * 10
                    )

        return matrix

    def run_simulation(self,
                       total_time: float = 10.0,
                       time_step: float = 0.1) -> List[Dict]:
        """运行混合流量模拟"""
        records = []

        while self.time < total_time:
            # 基础流量模式选择
            if np.random.rand() < 0.7: # 70%概率选择all_to_all
                base_matrix = self.generate_all_to_all()
                traffic_type = "all_to_all"
            else:
                base_matrix = self.generate_ring_allreduce()
                traffic_type = "ring_allreduce"

            # 突发流量叠加
            if np.random.rand() < 0.2:  # 20%概率触发突发
                burst_matrix = self.generate_burst_traffic()
                combined_matrix = base_matrix + burst_matrix
                traffic_type += "+burst"
            else:
                combined_matrix = base_matrix.copy()
                traffic_type += "+base"

            # 记录时间戳和矩阵
            records.append({
                "timestamp": round(self.time, 2),
                "type": traffic_type,
                "matrix": combined_matrix
            })

            # 更新时间（带随机波动）
            self.time += time_step * np.random.uniform(0.8, 1.2)

        return records


def analyze_traffic(records: List[Dict]):
    """增强型流量分析"""
    stats = defaultdict(float)
    hotspot_stats = defaultdict(float)

    for record in records:
        matrix = np.array(record["matrix"])
        stats[record["type"]] += matrix.sum()

        # 统计热点流量
        for src in TrafficGenerator().burst_config['hotspots']:
            hotspot_stats[src] += matrix[src, :].sum()

    print("\n=== 流量类型统计 ===")
    for t, total in stats.items():
        print(f"{t}: {total / 1024:.2f} GB")

    print("\n=== 热点节点统计 ===")
    for node, total in hotspot_stats.items():
        print(f"Node {node}: {total / 1024:.2f} GB")


if __name__ == "__main__":
    # 初始化带突发配置的生成器
    generator = TrafficGenerator(
        num_nodes=4,
        burst_config={
            'hotspots': [0, 1],
            'burst_duration': 0.3,
            'lambda_poisson': 0.4,
            'burst_intensity': 100.0
        }
    )

    # 运行10秒的模拟
    traffic_records = generator.run_simulation(total_time=10.0)

    # 保存结果
    with open("traffic_simulation.json", "w") as f:
        json.dump(traffic_records, f, indent=2, cls=NumpyEncoder)

    # 显示分析结果
    analyze_traffic(traffic_records)
    print("\n模拟数据已保存至 traffic_simulation.json")

    # 绘制每个时间点的流量矩阵热力图
    # for record in traffic_records:
    #     matrix = record["matrix"]
    #     timestamp = record["timestamp"]
    #     traffic_type = record["type"]
    #     plt.figure(figsize=(10, 8))
    #     sns.heatmap(matrix, annot=False, fmt=".2f", cmap="YlGnBu")
    #     plt.title(f'Traffic Matrix Heatmap at Time {timestamp} ({traffic_type})')
    #     plt.xlabel('Destination Node')
    #     plt.ylabel('Source Node')
    #     plt.show()
    #     plt.savefig("D:/桌面/之江工作/毕业成果/论文/论文实验/实验结果图（16个节点）/流量生成图2/temp{}.png".format(i))  # 输入地址，并利用format函数修改图片名称
    #     plt.clf()  # 需要重新更新画布，否则会出现同一张画布上绘制多张图片


    # # 绘制每个时间点的流量矩阵热力图
    # for i, record in enumerate(traffic_records):
    #     matrix = record["matrix"]
    #     timestamp = record["timestamp"]
    #     traffic_type = record["type"]
    #     plt.figure(figsize=(10, 8))
    #     sns.heatmap(matrix, annot=False, fmt=".2f", cmap="YlGnBu")
    #     plt.title(f'Traffic Matrix Heatmap at Time {timestamp} ({traffic_type})')
    #     plt.xlabel('Destination Node')
    #     plt.ylabel('Source Node')
    #     plt.savefig(
    #         f"D:/桌面/之江工作/毕业成果/论文/论文实验/实验结果图（16个节点）/流量生成图2/temp{i}.png")  # 使用f-string格式化字符串
    #     plt.clf()  # 需要重新更新画布，否则会出现同一张画布上绘制多张图片


    # stats = defaultdict(float)
    # for record in traffic_records:
    #     matrix = np.array(record["matrix"])
    #     stats[record["type"]] += matrix.sum()
    #
    # types = list(stats.keys())
    # totals = [stats[t] / 1024 for t in types]
    #
    # plt.figure(figsize=(10, 6))
    # plt.bar(types, totals)
    # plt.title('Total Traffic by Traffic Type')
    # plt.xlabel('Traffic Type')
    # plt.ylabel('Total Traffic (GB)')
    # plt.xticks(rotation=45)
    # plt.show()

    import matplotlib.pyplot as plt

    hotspot_nodes = [0, 1]
    hotspot_traffic_over_time = {node: [] for node in hotspot_nodes}
    timestamps = []

    for record in traffic_records:
        matrix = np.array(record["matrix"])
        timestamp = record["timestamp"]
        timestamps.append(timestamp)
        for node in hotspot_nodes:
            hotspot_traffic_over_time[node].append(matrix[node, :].sum())

    plt.figure(figsize=(10, 6))
    for node in hotspot_nodes:
        plt.plot(timestamps, hotspot_traffic_over_time[node], label=f'Node {node}')

    plt.title('Hotspot Node Traffic over Time')
    plt.xlabel('Time (s)')
    plt.ylabel('Traffic Volume')
    plt.legend()
    plt.show()


